#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copy-Move处理脚本
基于SAM模型对图像进行copy-move处理，并生成相应的mask和JSON元数据
重构版本：采用面向对象设计，实现新的copy-move逻辑
"""

import os
import cv2
import numpy as np
import torch
import json
import glob
import shutil
import re
import gc
import random
import logging
from datetime import datetime
from typing import Tuple, List, Dict, Any, Optional
from tqdm import tqdm
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry

# 配置常量
SAM_PATH = "/path/to/sam_vit_h_4b8939.pth"
LOG_DIR = "/path/to/logs"
DATASET_DIR = "/path/to/dataset"
RESULT_DIR = "/path/to/result/cpmv"


class ImageProcessor:
    """单个图像处理器"""

    def __init__(self, image_path: str, mask_generator: SamAutomaticMaskGenerator, logger: logging.Logger, seed: int = 42):
        """
        初始化图像处理器

        Args:
            image_path: 图像文件路径
            mask_generator: SAM掩码生成器
            logger: 日志记录器
            seed: 随机种子，确保可复现性
        """
        self.image_path = image_path
        self.mask_generator = mask_generator
        self.logger = logger
        self.seed = seed

        # 设置随机种子
        random.seed(seed)
        np.random.seed(seed)

        # 基本属性
        self.image: Optional[np.ndarray] = None
        self.image_shape: Optional[Tuple[int, int, int]] = None
        self.paper_id: Optional[str] = None
        self.filename: Optional[str] = None

        # 处理结果
        self.selected_object = None
        self.original_mask = None
        self.transparent_mask = None
        self.white_bg_mask = None
        self.copy_moved_image = None
        self.combined_mask = None
        self.target_position = None
        self.has_valid_result = False

        # 初始化
        self._load_and_validate()

    def _load_and_validate(self):
        """加载和验证图像"""
        # 提取paper_id
        self.paper_id = self._extract_paper_id()
        if not self.paper_id:
            raise ValueError(f"无法从文件名提取paper_id: {self.image_path}")

        # 获取文件名
        self.filename = os.path.basename(self.image_path)

        # 读取图像
        self.image = cv2.imread(self.image_path)
        if self.image is None:
            raise ValueError(f"无法读取图像: {self.image_path}")

        self.image_shape = self.image.shape
        self.logger.info(f"图像 {self.filename} 加载成功，尺寸: {self.image_shape}")

    def _extract_paper_id(self) -> Optional[str]:
        """从文件名中提取paper_id"""
        basename = os.path.basename(self.image_path)
        name_without_ext = os.path.splitext(basename)[0]

        # 使用正则表达式匹配 XX_数字_XX_XX 格式
        pattern = r'^\w+_(\d+)_\d+_\d+$'
        match = re.match(pattern, name_without_ext)

        if match:
            return match.group(1)
        return None

    def _clean_mask(self, mask: np.ndarray) -> np.ndarray:
        """清理mask，去除噪声"""
        # 转换为二值图像
        _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

        # 形态学操作：去除小噪声
        kernel = np.ones((3, 3), np.uint8)

        # 开运算：去除小的白色噪声
        cleaned_mask = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel, iterations=1)

        # 闭运算：填充小的黑色空洞
        cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel, iterations=1)

        # 高斯模糊平滑边缘
        cleaned_mask = cv2.GaussianBlur(cleaned_mask, (3, 3), 0)

        # 重新二值化
        _, cleaned_mask = cv2.threshold(cleaned_mask, 127, 255, cv2.THRESH_BINARY)

        return cleaned_mask

    def _extract_objects(self) -> List[Dict[str, Any]]:
        """提取符合条件的对象（面积比0.1-0.4）"""
        if self.image is None:
            return []

        try:
            # 生成masks
            masks = self.mask_generator.generate(self.image)

            if not masks:
                self.logger.warning(f"图像 {self.filename} 未检测到任何对象")
                return []

            # 按面积排序masks
            masks.sort(key=lambda x: x["area"], reverse=True)

            # 计算图像总面积
            if self.image_shape is None:
                return []
            image_area = self.image_shape[0] * self.image_shape[1]

            # 筛选对象 (0.1-0.4)
            valid_objects = []

            for mask_info in masks:
                area_ratio = mask_info["area"] / image_area

                if 0.1 <= area_ratio <= 0.4:
                    # 清理mask
                    raw_mask = mask_info["segmentation"].astype(np.uint8) * 255
                    cleaned_mask = self._clean_mask(raw_mask)

                    obj_info = {
                        "mask": cleaned_mask,
                        "area": mask_info["area"],
                        "area_ratio": area_ratio,
                        "bbox": mask_info.get("bbox", None)
                    }
                    valid_objects.append(obj_info)

            self.logger.info(f"图像 {self.filename} 检测到 {len(masks)} 个对象，符合条件: {len(valid_objects)} 个")
            return valid_objects

        except Exception as e:
            self.logger.error(f"提取对象时出错: {e}")
            return []

    def _select_object(self, objects: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """选择对象进行copy-move"""
        if not objects:
            return None

        # 随机选择一个对象
        selected = random.choice(objects)
        self.logger.info(f"选择对象进行copy-move，面积占比: {selected['area_ratio']:.2%}")
        return selected

    def _generate_masks(self, object_mask: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """生成两种类型的遮罩"""
        if self.image_shape is None:
            return np.array([]), np.array([])

        # 透明遮罩 (RGBA格式，选中区域为不透明白色，其他区域为透明)
        transparent_mask = np.zeros((*self.image_shape[:2], 4), dtype=np.uint8)
        mask_coords = np.where(object_mask > 0)
        transparent_mask[mask_coords[0], mask_coords[1]] = [255, 255, 255, 255]  # 白色不透明

        # 白底黑色遮罩 (选中区域为黑色，其他区域为白色)
        white_bg_mask = np.ones(self.image_shape[:2], dtype=np.uint8) * 255
        white_bg_mask[mask_coords[0], mask_coords[1]] = 0  # 黑色

        return transparent_mask, white_bg_mask

    def _scale_mask_if_needed(self, mask: np.ndarray, area_ratio: float) -> np.ndarray:
        """如果面积比>0.3则缩放mask到0.3以内"""
        if area_ratio <= 0.3:
            return mask

        # 计算缩放比例
        target_ratio = 0.3
        scale_factor = np.sqrt(target_ratio / area_ratio)

        self.logger.info(f"缩放mask，从 {area_ratio:.2%} 到 {target_ratio:.2%}，缩放因子: {scale_factor:.3f}")

        # 获取mask的非零区域
        mask_coords = np.where(mask > 0)
        if len(mask_coords[0]) == 0:
            return mask

        # 计算边界框
        min_y, max_y = mask_coords[0].min(), mask_coords[0].max()
        min_x, max_x = mask_coords[1].min(), mask_coords[1].max()

        # 计算中心点
        center_y = (min_y + max_y) // 2
        center_x = (min_x + max_x) // 2

        # 提取mask区域
        mask_region = mask[min_y:max_y+1, min_x:max_x+1]

        # 缩放mask区域
        new_h = int((max_y - min_y + 1) * scale_factor)
        new_w = int((max_x - min_x + 1) * scale_factor)

        # 确保最小尺寸
        new_h = max(new_h, 10)
        new_w = max(new_w, 10)

        scaled_region = cv2.resize(mask_region, (new_w, new_h), interpolation=cv2.INTER_NEAREST)

        # 创建新的mask
        scaled_mask = np.zeros_like(mask)

        # 计算放置位置（以中心对齐）
        start_y = max(0, center_y - new_h // 2)
        end_y = min(mask.shape[0], start_y + new_h)
        start_x = max(0, center_x - new_w // 2)
        end_x = min(mask.shape[1], start_x + new_w)

        # 调整实际尺寸
        actual_h = end_y - start_y
        actual_w = end_x - start_x

        if actual_h > 0 and actual_w > 0:
            if actual_h != new_h or actual_w != new_w:
                scaled_region = cv2.resize(scaled_region, (actual_w, actual_h), interpolation=cv2.INTER_NEAREST)
            scaled_mask[start_y:end_y, start_x:end_x] = scaled_region

        return scaled_mask

    def _find_target_position_with_noise(self, original_mask: np.ndarray) -> Tuple[int, int]:
        """找到最远位置并添加噪声偏置"""
        if self.image_shape is None:
            return 50, 50
        img_h, img_w = self.image_shape[:2]

        # 获取原始对象的中心位置
        mask_coords = np.where(original_mask > 0)
        if len(mask_coords[0]) == 0:
            return img_w // 2, img_h // 2

        orig_center_y = int(np.mean(mask_coords[0]))
        orig_center_x = int(np.mean(mask_coords[1]))

        # 计算对象边界框
        min_y, max_y = mask_coords[0].min(), mask_coords[0].max()
        min_x, max_x = mask_coords[1].min(), mask_coords[1].max()
        obj_h, obj_w = max_y - min_y + 1, max_x - min_x + 1

        # 生成候选位置网格
        step = max(20, min(obj_w, obj_h) // 2)
        candidates = []

        for y in range(obj_h // 2, img_h - obj_h // 2, step):
            for x in range(obj_w // 2, img_w - obj_w // 2, step):
                # 计算距离原始中心的距离
                distance = np.sqrt((x - orig_center_x)**2 + (y - orig_center_y)**2)
                candidates.append((x, y, distance))

        if not candidates:
            # 如果没有合适的候选位置，返回图像中心
            return img_w // 2, img_h // 2

        # 选择距离最远的位置
        candidates.sort(key=lambda x: x[2], reverse=True)
        best_x, best_y, best_distance = candidates[0]

        # 添加噪声偏置（使用固定seed确保可复现）
        noise_range = min(50, obj_w // 4, obj_h // 4)  # 噪声范围
        noise_x = random.randint(-noise_range, noise_range)
        noise_y = random.randint(-noise_range, noise_range)

        final_x = np.clip(best_x + noise_x, obj_w // 2, img_w - obj_w // 2)
        final_y = np.clip(best_y + noise_y, obj_h // 2, img_h - obj_h // 2)

        self.logger.info(f"目标位置: ({final_x}, {final_y}), 原始中心: ({orig_center_x}, {orig_center_y}), 噪声偏置: ({noise_x}, {noise_y})")
        return final_x, final_y

    def _clip_mask_to_bounds(self, mask: np.ndarray, center_x: int, center_y: int) -> Tuple[np.ndarray, int, int]:
        """将遮罩裁切到图像边界内，并返回裁切后的遮罩和实际放置位置"""
        if self.image_shape is None:
            return mask, center_x, center_y
        img_h, img_w = self.image_shape[:2]

        # 获取mask的边界框
        mask_coords = np.where(mask > 0)
        if len(mask_coords[0]) == 0:
            return mask, center_x, center_y

        min_y, max_y = mask_coords[0].min(), mask_coords[0].max()
        min_x, max_x = mask_coords[1].min(), mask_coords[1].max()

        # 计算mask的中心偏移量
        mask_center_y = (min_y + max_y) // 2
        mask_center_x = (min_x + max_x) // 2

        # 计算新位置的边界框
        offset_y = center_y - mask_center_y
        offset_x = center_x - mask_center_x

        new_min_y = min_y + offset_y
        new_max_y = max_y + offset_y
        new_min_x = min_x + offset_x
        new_max_x = max_x + offset_x

        # 计算需要裁切的部分
        clip_top = max(0, -new_min_y)
        clip_bottom = max(0, new_max_y - (img_h - 1))
        clip_left = max(0, -new_min_x)
        clip_right = max(0, new_max_x - (img_w - 1))

        # 如果需要裁切
        if clip_top > 0 or clip_bottom > 0 or clip_left > 0 or clip_right > 0:
            self.logger.info(f"需要裁切mask，裁切量: top={clip_top}, bottom={clip_bottom}, left={clip_left}, right={clip_right}")

            # 创建裁切后的mask
            clipped_mask = np.zeros_like(mask)

            # 重新计算边界
            src_top = min_y + clip_top
            src_bottom = max_y - clip_bottom + 1
            src_left = min_x + clip_left
            src_right = max_x - clip_right + 1

            dst_top = max(0, new_min_y)
            dst_bottom = min(img_h, new_max_y + 1)
            dst_left = max(0, new_min_x)
            dst_right = min(img_w, new_max_x + 1)

            if src_top < src_bottom and src_left < src_right and dst_top < dst_bottom and dst_left < dst_right:
                # 复制有效区域
                src_region = mask[src_top:src_bottom, src_left:src_right]
                clipped_mask[dst_top:dst_bottom, dst_left:dst_right] = src_region

            # 更新实际放置位置
            actual_center_x = (dst_left + dst_right) // 2
            actual_center_y = (dst_top + dst_bottom) // 2

            return clipped_mask, actual_center_x, actual_center_y
        else:
            # 不需要裁切，直接移动mask
            moved_mask = np.zeros_like(mask)
            moved_mask[new_min_y:new_max_y+1, new_min_x:new_max_x+1] = mask[min_y:max_y+1, min_x:max_x+1]
            return moved_mask, center_x, center_y

    def _perform_copy_move(self, original_mask: np.ndarray, target_x: int, target_y: int) -> Tuple[np.ndarray, np.ndarray]:
        """执行copy-move操作"""
        if self.image is None or self.image_shape is None:
            return np.array([]), np.array([])

        # 裁切遮罩到边界内
        clipped_mask, actual_x, actual_y = self._clip_mask_to_bounds(original_mask, target_x, target_y)

        # 创建copy-moved图像
        copy_moved_image = self.image.copy()

        # 获取原始mask的像素
        original_coords = np.where(original_mask > 0)
        clipped_coords = np.where(clipped_mask > 0)

        if len(original_coords[0]) > 0 and len(clipped_coords[0]) > 0:
            # 执行像素复制
            for i, (orig_y, orig_x) in enumerate(zip(original_coords[0], original_coords[1])):
                if i < len(clipped_coords[0]):
                    new_y, new_x = clipped_coords[0][i], clipped_coords[1][i]
                    copy_moved_image[new_y, new_x] = self.image[orig_y, orig_x]

        # 创建组合mask（原始mask + 移动后的mask）
        combined_mask = np.zeros(self.image_shape[:2], dtype=np.uint8)
        combined_mask[original_mask > 0] = 255  # 原始位置
        combined_mask[clipped_mask > 0] = 255   # 新位置

        return copy_moved_image, combined_mask

    def process(self) -> bool:
        """处理图像，执行完整的copy-move流程"""
        try:
            self.logger.info(f"开始处理图像: {self.filename}")

            # 1. 提取对象
            objects = self._extract_objects()
            if not objects:
                self.logger.warning(f"图像 {self.filename} 没有符合条件的对象")
                return False

            # 2. 选择对象
            self.selected_object = self._select_object(objects)
            if not self.selected_object:
                return False

            # 3. 生成两种遮罩
            self.transparent_mask, self.white_bg_mask = self._generate_masks(self.selected_object["mask"])

            # 4. 缩放遮罩（如果需要）
            scaled_mask = self._scale_mask_if_needed(
                self.selected_object["mask"],
                self.selected_object["area_ratio"]
            )

            # 5. 找到目标位置（最远点 + 噪声）
            target_x, target_y = self._find_target_position_with_noise(scaled_mask)
            self.target_position = (target_x, target_y)

            # 6. 执行copy-move
            self.copy_moved_image, self.combined_mask = self._perform_copy_move(scaled_mask, target_x, target_y)

            # 7. 验证结果
            if np.any(self.combined_mask > 0):
                self.has_valid_result = True
                self.logger.info(f"图像 {self.filename} 处理成功")
                return True
            else:
                self.logger.warning(f"图像 {self.filename} 没有生成有效的mask")
                return False

        except Exception as e:
            self.logger.error(f"处理图像 {self.filename} 时出错: {e}")
            return False

    def get_result_info(self) -> Dict[str, Any]:
        """获取处理结果信息"""
        if self.filename is None:
            panel_id = "unknown"
        else:
            panel_id = os.path.splitext(self.filename)[0]

        return {
            "paper_id": f"{self.paper_id}_FG_002",
            "panel_id": panel_id,
            "filename": self.filename,
            "has_valid_result": self.has_valid_result,
            "original_path": self.image_path,
            "target_position": self.target_position,
            "area_ratio": self.selected_object["area_ratio"] if self.selected_object else 0
        }


class CopyMoveHandler:
    """Copy-Move主处理器"""

    def __init__(self, seed: int = 42):
        """
        初始化处理器

        Args:
            seed: 随机种子，确保可复现性
        """
        self.sam_path = SAM_PATH
        self.log_dir = LOG_DIR
        self.dataset_dir = DATASET_DIR
        self.result_dir = RESULT_DIR
        self.seed = seed

        # 设置全局随机种子
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.sam_model = None
        self.mask_generator = None
        self.logger = None

        # 输出目录
        self.masked_dir = os.path.join(self.result_dir, "masked")
        self.nomasked_dir = os.path.join(self.result_dir, "nomasked")

        # 处理结果
        self.masked_results = []
        self.nomasked_results = []

        # 初始化
        self._setup_directories()
        self._setup_logging()

    def _setup_directories(self):
        """设置目录结构"""
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.masked_dir, exist_ok=True)
        os.makedirs(self.nomasked_dir, exist_ok=True)

    def _setup_logging(self):
        """设置日志系统"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = os.path.join(self.log_dir, f"cpmv-{timestamp}.log")

        # 创建logger
        self.logger = logging.getLogger('CopyMoveHandler')
        self.logger.setLevel(logging.INFO)

        # 清除现有handlers
        if self.logger.handlers:
            self.logger.handlers.clear()

        # 创建文件handler
        file_handler = logging.FileHandler(log_file, encoding='utf-8')
        file_handler.setLevel(logging.INFO)

        # 创建控制台handler
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)

        # 创建formatter
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)

        # 添加handlers
        self.logger.addHandler(file_handler)
        # self.logger.addHandler(console_handler)

        self.logger.info(f"日志系统初始化完成，日志文件: {log_file}")

    def _load_sam_model(self):
        """加载SAM模型"""
        if self.sam_model is None:
            if self.logger:
                self.logger.info("开始加载SAM模型...")
                self.logger.info(f"模型路径: {self.sam_path}")
                self.logger.info(f"计算设备: {self.device}")

            model_type = "vit_h"
            self.sam_model = sam_model_registry[model_type](checkpoint=self.sam_path)
            self.sam_model.to(device=self.device)

            # 配置掩码生成器
            self.mask_generator = SamAutomaticMaskGenerator(
                self.sam_model,
                output_mode="binary_mask",
                points_per_side=16,
                pred_iou_thresh=0.86,
                stability_score_thresh=0.92,
                min_mask_region_area=100
            )
            if self.logger:
                self.logger.info("SAM模型加载完成!")

    def _scan_dataset(self) -> List[str]:
        """扫描数据集目录"""
        if self.logger:
            self.logger.info(f"扫描数据集目录: {self.dataset_dir}")

        image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff']
        image_files = []

        for ext in image_extensions:
            pattern = os.path.join(self.dataset_dir, ext)
            files = glob.glob(pattern)
            image_files.extend(files)

        if self.logger:
            self.logger.info(f"找到 {len(image_files)} 个图像文件")
        return image_files

    def _create_output_structure(self, paper_id: str, has_mask: bool) -> Dict[str, str]:
        """创建输出目录结构"""
        fg_paper_id = f"{paper_id}_FG_002"

        if has_mask:
            base_dir = os.path.join(self.masked_dir, fg_paper_id)
        else:
            base_dir = os.path.join(self.nomasked_dir, fg_paper_id)

        directories = {
            'base': base_dir,
            'figure': os.path.join(base_dir, 'figure'),
            'mask': os.path.join(base_dir, 'mask'),
            'panel': os.path.join(base_dir, 'panel')
        }

        # 创建目录
        for dir_path in directories.values():
            os.makedirs(dir_path, exist_ok=True)

        return directories

    def _save_image_results(self, processor: ImageProcessor) -> Optional[Dict[str, Any]]:
        """保存图像处理结果"""
        result_info = processor.get_result_info()

        if processor.paper_id is None:
            if self.logger:
                self.logger.error("paper_id为空，跳过保存")
            return None

        # 创建输出目录
        directories = self._create_output_structure(
            processor.paper_id,
            processor.has_valid_result
        )

        try:
            # 1. 保存原图到figure目录
            if processor.filename is None:
                if self.logger:
                    self.logger.error("filename为空，跳过保存")
                return None

            figure_path = os.path.join(directories['figure'], processor.filename)
            shutil.copy2(processor.image_path, figure_path)

            if processor.has_valid_result:
                # 2. 保存处理后的图像到panel目录
                panel_path = os.path.join(directories['panel'], processor.filename)
                if processor.copy_moved_image is not None:
                    cv2.imwrite(panel_path, processor.copy_moved_image)

                # 3. 保存组合mask到mask目录
                mask_path = os.path.join(directories['mask'], processor.filename)
                # 保存为白底黑色mask（反转）
                if processor.combined_mask is not None:
                    cv2.imwrite(mask_path, 255 - processor.combined_mask)

                # 构建结果信息
                json_entry = {
                    "question_id": f"cpmv_{result_info['panel_id']}",
                    "issue": ["Image forgery"],
                    "sub_issue": ["copy-move"],
                    "answer": ["B"],
                    "paper_id": result_info["paper_id"],
                    "paper_path": None,
                    "panel_id": result_info["panel_id"],
                    "panel_path": [os.path.abspath(panel_path)],
                    "has_issue": True,
                    "mask": [],
                    "mask_path": [os.path.abspath(mask_path)],
                    "mask_caption": None
                }

                self.masked_results.append(json_entry)
                if self.logger:
                    self.logger.info(f"保存masked结果: {processor.filename}")

            else:
                # 构建无mask的结果信息
                json_entry = {
                    "question_id": f"cpmv_{result_info['panel_id']}",
                    "issue": ["Image forgery"],
                    "sub_issue": ["copy-move"],
                    "answer": ["A"],
                    "paper_id": result_info["paper_id"],
                    "paper_path": None,
                    "panel_id": result_info["panel_id"],
                    "panel_path": [os.path.abspath(figure_path)],
                    "has_issue": False,
                    "mask": [],
                    "mask_path": [],
                    "mask_caption": None
                }

                self.nomasked_results.append(json_entry)
                if self.logger:
                    self.logger.info(f"保存nomasked结果: {processor.filename}")

            return json_entry

        except Exception as e:
            if self.logger:
                self.logger.error(f"保存结果时出错: {e}")
            return None

    def _generate_json_files(self):
        """生成JSON文件"""
        if self.logger:
            self.logger.info("生成JSON元数据文件...")

        # 1. 保存masked.json
        if self.masked_results:
            masked_json_path = os.path.join(self.masked_dir, "masked.json")
            with open(masked_json_path, 'w', encoding='utf-8') as f:
                json.dump(self.masked_results, f, indent=4, ensure_ascii=False)
            if self.logger:
                self.logger.info(f"Masked JSON文件已保存: {masked_json_path} ({len(self.masked_results)} 条记录)")

        # 2. 保存nomasked.json
        if self.nomasked_results:
            nomasked_json_path = os.path.join(self.nomasked_dir, "nomasked.json")
            with open(nomasked_json_path, 'w', encoding='utf-8') as f:
                json.dump(self.nomasked_results, f, indent=4, ensure_ascii=False)
            if self.logger:
                self.logger.info(f"NoMasked JSON文件已保存: {nomasked_json_path} ({len(self.nomasked_results)} 条记录)")

        # 3. 保存总的JSON文件
        all_results = self.masked_results + self.nomasked_results
        if all_results:
            all_json_path = os.path.join(self.result_dir, "cpmv_all.json")
            with open(all_json_path, 'w', encoding='utf-8') as f:
                json.dump(all_results, f, indent=4, ensure_ascii=False)
            if self.logger:
                self.logger.info(f"总JSON文件已保存: {all_json_path} ({len(all_results)} 条记录)")

        # 4. 生成统计信息JSON
        stats = {
            "total_images": len(all_results),
            "masked_count": len(self.masked_results),
            "nomasked_count": len(self.nomasked_results),
            "masked_percentage": len(self.masked_results) / len(all_results) * 100 if all_results else 0,
            "processing_time": datetime.now().isoformat(),
            "files": {
                "masked_json": "masked/masked.json" if self.masked_results else None,
                "nomasked_json": "nomasked/nomasked.json" if self.nomasked_results else None,
                "all_json": "cpmv_all.json"
            }
        }

        stats_path = os.path.join(self.result_dir, "processing_stats.json")
        with open(stats_path, 'w', encoding='utf-8') as f:
            json.dump(stats, f, indent=4, ensure_ascii=False)
        if self.logger:
            self.logger.info(f"统计信息已保存: {stats_path}")

    def process_all_images(self):
        """处理所有图像"""
        if self.logger:
            self.logger.info("开始批量处理图像")

        # 扫描数据集
        image_files = self._scan_dataset()
        if not image_files:
            if self.logger:
                self.logger.warning("未找到任何图像文件")
            return

        # 加载SAM模型
        self._load_sam_model()

        # 处理每个图像
        for image_path in tqdm(image_files, desc="处理图像"):
            try:
                # 创建图像处理器
                if self.mask_generator is None or self.logger is None:
                    print("错误：mask_generator或logger未初始化")
                    continue

                processor = ImageProcessor(
                    image_path=image_path,
                    mask_generator=self.mask_generator,
                    logger=self.logger,
                    seed=self.seed
                )

                # 处理图像
                processor.process()

                # 保存结果
                self._save_image_results(processor)

                # 清理内存
                del processor
                torch.cuda.empty_cache()
                gc.collect()

            except Exception as e:
                if self.logger:
                    self.logger.error(f"处理图像 {image_path} 时出错: {e}")
                continue

        if self.logger:
            self.logger.info(f"处理完成，成功处理 {len(self.masked_results) + len(self.nomasked_results)} 个图像")

    def run(self):
        """运行完整的处理流程"""
        try:
            if self.logger:
                self.logger.info("开始Copy-Move处理流程")

            # 处理所有图像
            self.process_all_images()

            # 生成JSON文件
            self._generate_json_files()

            if self.logger:
                self.logger.info("Copy-Move处理流程完成")

        except Exception as e:
            if self.logger:
                self.logger.error(f"处理流程出错: {e}")
            raise
        finally:
            # 清理资源
            if self.logger:
                self.logger.info("清理资源...")
            if self.sam_model is not None:
                del self.sam_model
            torch.cuda.empty_cache()
            gc.collect()
            if self.logger:
                self.logger.info("资源清理完成")


def main():
    """主函数"""
    print("启动Copy-Move处理程序 (重构版)")

    # 创建处理器并运行
    handler = CopyMoveHandler(seed=42)  # 使用固定种子确保可复现性
    handler.run()

    print("程序执行完成")


if __name__ == "__main__":
    main()
